from flask import Flask, request, jsonify
from flask_cors import CORS
import json
import llmShift

app = Flask(__name__)
CORS(app)
CORS(app, resources={r"/api/*": {"origins": "http://localhost:63342"}})
model = list(llmShift.llm_detail.keys())
chosen_model = []
for key in model:
    chosen_model.append(list(list(llmShift.llm_detail[key]["model_name"].keys())))

def json_generate(model,chosen_model):
    json_data = []
    for i, label in enumerate(model):
        children = [{'value': str(chosen_model), 'label': str(chosen_model)} for chosen_model in chosen_model[i]]
        json_data.append({
            'value': label,
            'label': label,
            'children': children
        })
    # 将Python对象转换为JSON字符串
    json_str = json.dumps(json_data, indent=2, ensure_ascii=False)
    return json_str
#print(json_generate(model,chosen_model))


def arraysSplit(chosen_model,temperature,question):
    answer = []
    print(chosen_model)
    for i in range(len(chosen_model)):
        answer.append(llmShift.version_choose(chosen_model[i][1],temperature,question))
        print(llmShift.version_choose(chosen_model[i][1],temperature,question))
    return answer


@app.route('/get_info', methods=['POST'])
def get_info():
    return jsonify({
        "code": 200,
        "msg": "success",
        "data": json_generate(model,chosen_model)
    })


@app.route('/qa', methods=['POST'])
def qa():
    data = request.get_json()
    print(data)
    return jsonify({
        "code": 200,
        "msg": "success",
        "data": arraysSplit(data.get('chosen_model'), data.get('temperature'), data.get('question'))
    })


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7878, debug=True)
